// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: jhkim

#pragma once

#include <set>
#include "presto/quant/base/logging.h"
#include "presto/quant/math/moving_window.h"

namespace presto::math {

template <typename Elem>
class TimeWindowMedian {
 public:
  TimeWindowMedian() : moving_window_(this) {
    ResizeWindow(1);
    num_popped = 0;
  }

  explicit TimeWindowMedian(int64_t window_size) : moving_window_(this) {
    ResizeWindow(window_size);
    num_popped = 0;
  }

  ~TimeWindowMedian() = default;

  void ResizeWindow(int64_t window_size) {
    CHECK_GT(window_size, 0);
    moving_window_.ResizeWindow(window_size);
    Rebalance();
  }

  void clear() {
    ResizeWindow(moving_window_.window_size());
  }

  void Update(int64_t time, const Elem& element) {
    moving_window_.Update(time, element);
    Rebalance();
  }

  void UpdateTimeOnly(int64_t time) {
    moving_window_.UpdateTimeOnly(time);
    Rebalance();
  }

  const Elem& GetMedian() const {
    return last_median_;
  }

  int size() const {
    return low_set_.size() + high_set_.size();
  }

  // Internal use only.
  void OnPushed(int64_t /*time*/, const Elem& elem) {
    if (low_set_.empty() || *low_set_.rbegin() >= elem) {
      low_set_.insert(elem);
    } else {
      high_set_.insert(elem);
    }
  }

  void OnPopped(int64_t /*time*/, const Elem& elem) {
    num_popped += 1;
    DCHECK_GT(size(), 0);
    if (!low_set_.empty() && *low_set_.rbegin() >= elem) {
      typename std::multiset<Elem>::iterator it = low_set_.find(elem);
      DCHECK(it != low_set_.end());
      low_set_.erase(it);
    } else {
      typename std::multiset<Elem>::iterator it = high_set_.find(elem);
      DCHECK(it != high_set_.end());
      high_set_.erase(it);
    }
  }

  int64_t num_popped;

 private:
  MovingWindowWithCallee<Elem, TimeWindowMedian<Elem>> moving_window_;

  void ResetMedianValue() {
    if (low_set_.size() > 0) {
      last_median_ = *low_set_.rbegin();
    }
  }

  // We maintain the size of low and high set.
  //   low_set_.size() == high_set_.size() or
  //   low_set_.size() == higt_set_.size() + 1
  void Rebalance() {
    while (low_set_.size() > high_set_.size() + 1) {
      typename std::multiset<Elem>::iterator it = low_set_.end();
      --it;

      high_set_.insert(*it);
      low_set_.erase(it);
    }

    while (low_set_.size() < high_set_.size()) {
      typename std::multiset<Elem>::iterator it = high_set_.begin();
      low_set_.insert(*it);
      high_set_.erase(it);
    }

    ResetMedianValue();
  }

  Elem last_median_;

  // We define the median as the highest value of low set.
  std::multiset<Elem> low_set_;
  std::multiset<Elem> high_set_;
};  // class TimeWindowMedian

}  // namespace presto::math
